import os
import time
import torch
import torchvision
import pandas as pd
import torch.nn as nn
from utils import progress_bar, get_data
import torchvision.transforms as transforms

# 解决随机性而产生的结果有稍微不同，不过可能会因为Pytorch版本不同而有一定差别
SEED = 1
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

train_csv = pd.DataFrame(columns=['Epoch', 'StepTime', 'TotalTime', 'PerSec', 'Acc', 'Total'])

BATCH_SIZE = 32
EPOCHS = 5
criterion = nn.CrossEntropyLoss()

# 学习率优化
def adjust_learning_rate(optimizer, epoch, lr):
    lr *= (0.1 ** (epoch // 2))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

if __name__ == "__main__":

    # 网络定义
    device = "cpu"
    net = torchvision.models.resnet50(num_classes=10)
    net = net.to(device=device)

    # 数据载入增强
    trainset = torchvision.datasets.CIFAR10(
        root="./data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                ),
            ]
        ),
    )
    train_loader = torch.utils.data.DataLoader(
        trainset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
    )

    # 优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        net.parameters(),
        lr=0.01,
        momentum=0.9,
        weight_decay=0.0001,
        nesterov=True,
    )

    lr_init = optimizer.param_groups[0]['lr']

    print("            =======  Training  ======= \n")

    # 开始训练
    net.train()
    for ep in range(1, EPOCHS + 1):
        adjust_learning_rate(optimizer, ep, lr_init) # 可变学习率
        print("EPOCHS ", ep)
        train_loss = correct = total = 0

        for idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)

            loss = criterion(outputs, targets) #
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            total += targets.size(0)
            correct += torch.eq(outputs.argmax(dim=1), targets).sum().item()

            progress_bar(idx, len(train_loader), total, 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                         % (train_loss / (idx + 1), 100. * correct / total, correct, total))

        # 数据保存
        step_time, total_time, per_sec = get_data()
        train_csv = train_csv.append({'Epoch': ep, 'StepTime': step_time, 'TotalTime': total_time, 'PerSec': per_sec, 'Acc': 100. * correct / total, 'Total': total}, ignore_index=True)

    # 数据
    csv_path = os.getcwd() + "/csv/"
    now = time.strftime("%Y-%m-%d-%H_%M_%S", time.localtime(time.time()))
    save_name = csv_path + now + r"--ResNet50-6CPU.csv"
    train_csv.to_csv(save_name, mode='a', header=True, index=True)
    print("\n            =======  Training Finished  ======= \n")